package edu.kit.csl.pisa.models;

/*
This file is part of the PISA Alignment Tool.

Copyright (C) 2013
Karlsruhe Institute of Technology
Cognitive Systems Lab (CSL)
Felix Stahlberg

PISA is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

PISA is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with PISA. If not, see <http://www.gnu.org/licenses/>.
*/

import java.util.Arrays;

import edu.kit.csl.pisa.datatypes.AlignmentPosition;
import edu.kit.csl.pisa.datatypes.Corpus;
import edu.kit.csl.pisa.datatypes.Dictionary;
import edu.kit.csl.pisa.datatypes.SentenceAlignment;
import edu.kit.csl.pisa.io.Logger;
import edu.kit.csl.pisa.ui.Configuration;

/**
 * This class implements the Model 3-phone alignment model. In-depth discussion
 * of this model can be found in the technical paper. The parameters are stored
 * as n-dimensional double arrays. The order in which the dimensions are
 * assigned is as follows:<br>
 * At first, the parameters on the right side of the "|" in the order in which 
 * these dependencies are denoted in the technical paper. The last dimension
 * will always store the left side of the "|" (in P(.|.) ):<br>
 * d:	1. Number of words in the English sentence (l)<br>
 * 		2. Number of words in the French sentence (k)<br>
 * 		3. English word position (A_0j)<br>
 * 		4. French word position (A_1j)<br>
 * 		(A_0j is moved back for cache-efficiency when calculating P(a,f|e))<br>
 * o:	1. English Word (e_A_0j)<br>
 * 		2. Relative French word position (A_2j)<br>
 * 		3. French word length in phonemes (A_3j)<br>
 * 
 * @see ModelP
 */
public class Model3P extends ModelP {
	
	/*
	 * Model parameter. Which dimension contains which values is denoted in the
	 * head of this class.
	 */
	protected double[][][] d;
	protected double[][] o;
	
	// Model instance for collecting fractional counts
	protected Model3P newModel;
	
	/**
	 * Sole constructor.
	 * 
	 * @param name name of the model
	 */
	public Model3P(String name) {
		super(name);
	}

	/**
	 * Calculate the log probability P(a,f|e). The formula can be found in the
	 * technical paper. It uses a simple distortion and word length model, a 
	 * rather complex model for lexical translation and fertilities similar to
	 * IBM Model 3. The formula is derived from the IBM Model 3 formula.  
	 * 
	 * @return the logarithm of P(a,f|e)
	 */
	@Override
	public double calculateAlignmentProbability(SentenceAlignment a) {
		final int srcLength = a.getSourceSentence().getLength(); // l
		final int nTrgtWords = a.getAlignment().countTargetWords(); // k
		double[] summands = new double[3 + 2 * srcLength + 2 * nTrgtWords
		                               + a.getTargetSentence().getLength()];
		// p0, p1, fertility, lexical translation
		int i = calculateModelPAlignmentProbability(summands, a);
		// Word length and distortion
		for (AlignmentPosition pos : a.eachTargetWord()) {
			if (pos.getSourceWordPosition() == 0) {
				continue;
			}
			summands[i++] = d[srcLength]
							 [pos.getSourceWordPosition()]
							 [pos.getTargetWordPosition()];
			summands[i++] = o[pos.getSourceWord()]
							 [pos.getTargetWordLength()];
		}
		return sum(summands, 0, i);
	}

	/* (non-Javadoc)
	 * @see edu.kit.csl.pisa.models.AlignmentModel#initializeFractionalCounts()
	 */
	@Override
	public void initializeFractionalCounts() {
		newModel = new Model3P(getName());
		newModel.allocateMemory(false);
	}
	
	/* (non-Javadoc)
	 * @see edu.kit.csl.pisa.models.AlignmentModel#writeBackFractionalCounts()
	 */
	@Override
	public void writeBackFractionalCounts() {
		normalizeParameters();
		this.n = newModel.n;
		this.t = newModel.t;
		this.d = newModel.d;
		this.o = newModel.o;
		this.p0 = newModel.p0;
		this.p1 = newModel.p1;
		newModel = null;
	}
	
	/**
	 * This template method (GoF) is the interface for implementing subclasses
	 * to multi-threaded collection of fractional counts.
	 * TODO: Use recursive D&C summation
	 * 
	 * @param pair sentence alignment
	 * @param weight weight of the alignment
	 * @see #collectFractionalCounts(Corpus)
	 */
	@Override
	public void aggregateFractionalCount(SentenceAlignment a, double weight) {
		aggregateModelPFractionalCount(newModel, a, weight);
		int nSrcWords = a.getSourceSentence().getLength();
		synchronized (newModel.d) {
			for (AlignmentPosition pos : a.eachTargetWord()) {
				if (pos.getSourceWordPosition() == 0) {
					continue;
				}
				newModel.d	[nSrcWords]
							[pos.getSourceWordPosition()]
							[pos.getTargetWordPosition()] += weight;
				newModel.o	[pos.getSourceWord()]
							[pos.getTargetWordLength()] += weight;
			}
		}
	}
	
	public double[][] getO() {
		return o;
	}
	
	/*
	 * This is a helper method of {@link #writeBackFractionalCounts()}.
	 * Normalize parameters in a separate model instance. Thereby we assume,
	 * that the parameters are containing fractional counts. This method
	 * realizes a limited amount of parallelism, in which each parameter is
	 * normalized in a separate thread.
	 */
	private void normalizeParameters() {
		Thread[] workers = new Thread[4];
		workers[0] = new Thread() {
			public void run() {
				for (double[] a1 : newModel.n) {
					normalize(a1);
				}
			}
		};
		workers[1] = new Thread() {
			public void run() {
				for (double[] a1 : newModel.o) {
					normalize(a1);
				}
			}
		};
		workers[2] = new Thread() {
			public void run() {
				//for (double[][][] a1 : newModel.d) {
					for (double[][] a2 : newModel.d) {
						for (double[] a3 : a2) {
							normalize(a3);
						}
					}
				//}
			}
		};
		workers[3] = new Thread() {
			public void run() {
				for (double[][] a1 : newModel.t) {
					for (double[] a2 : a1) {
						normalize(a2);
					}
				}
			}
		};
		for (int i = 0; i < 4; i++) {
			workers[i].start();
		}
		for (int i = 0; i < 4; i++) {
			try {
				workers[i].join();
			} catch (InterruptedException e) {
				Logger.getSingleton().error("Join of worker threads in EM "
						+ "training interrupted: " + e.getMessage());
			}
		}
	}
	
	/**
	 * Import model parameters from IBM Model 3 ({@link Model3}). Additional
	 * dependencies in the lexical translation table (t) are distributed
	 * uniformly.
	 * 
	 * @param srcModel should be a {@link Model3} instance
	 * @throws IllegalArgumentException if srcModel is not a 
	 * 				{@link Model3} instance
	 */
	@Override
	public void importModelParameters(AlignmentModel srcModel)
			throws IllegalArgumentException {
		if (srcModel instanceof Model3) {
			importParametersFromModel3((Model3) srcModel);
		} else if (srcModel instanceof SimpleModel) {
			distributeUniformly();
		} else {
			throw new IllegalArgumentException("Import is not implemented"
				+ " for this alignment model.");
		}
	}
	
	/*
	 * Distribute all parameters uniformly.
	 */
	private void distributeUniformly() {
		final int trgtDictSize = 
				Dictionary.getSingleton(Dictionary.TRGT_DICT).size();
		final int maxTrgtSenLength = 
				Configuration.getSingleton().getInteger("maxTrgtSenLen");
		allocateMemory(true);
		double prop = Math.log(1f / maxFertility);
		for (double[] a1 : n) {
			Arrays.fill(a1, prop);
		}
		prop = Math.log(1f / trgtDictSize);
		for (double[][] a1 : t) {
			for (double[] a2 : a1) {
				Arrays.fill(a2, prop);
			}
		}
		prop = Math.log(1f / maxWordLength);
		for (double[] a1 : o) {
				Arrays.fill(a1, prop);
		}
		prop = Math.log(1f / maxTrgtSenLength);
		//for (double[][][] a1 : d) {
			for (double[][] a2 : d) {
				for (double[] a3 : a2) {
					Arrays.fill(a3, prop);
				}
			}
		//}
	}
	
	/*
	 * Import parameters from IBM Model 3
	 * 
	 * @param model3 the source model
	 */
	private void importParametersFromModel3(Model3 model3) {
		final int srcDictSize = 
				Dictionary.getSingleton(Dictionary.SRC_DICT).size();
		final int maxTrgtSenLength = 
				Configuration.getSingleton().getInteger("maxTrgtSenLen");
		allocateMemory(true);
		setP1(model3.getParameterP1()); // Very roughly, actually higher
		// Transfer t parameters as good as possible
		double[][] lexicalTranslation = model3.getParameterT();
		for (int srcToken = 0; srcToken < srcDictSize; srcToken++) {
			double[] distrib = lexicalTranslation[srcToken];
			for (int pos = 0; pos < maxWordLength; pos++) {
				t[srcToken][pos] = distrib.clone();
			}
		}
		/* How to distribute uniformly:
		 * all probabilities with the same right-hand side (right from |) have 
		 * to sum up to 1.
		 */
		double prop = Math.log(1f / maxFertility);
		for (double[] a1 : n) {
			Arrays.fill(a1, prop);
		}
		prop = Math.log(1f / maxWordLength);
		for (double[] a1 : o) {
				Arrays.fill(a1, prop);
		}
		prop = Math.log(1f / maxTrgtSenLength);
		//for (double[][][] a1 : d) {
			for (double[][] a2 : d) {
				for (double[] a3 : a2) {
					Arrays.fill(a3, prop);
				}
			}
		//}
	}
	
	/*
	 * Allocate memory for model parameters.
	 * 
	 * @param print Set to TRUE to print parameter sizes on the screen
	 */
	private void allocateMemory(boolean print) {
		allocateModelPMemory(print);
		final int srcDictSize = 
				Dictionary.getSingleton(Dictionary.SRC_DICT).size() + 1;
		final int maxSrcSenLength = 
				Configuration.getSingleton().getInteger("maxSrcSenLen") + 1;
		final int maxTrgtSenLength = 
				Configuration.getSingleton().getInteger("maxTrgtSenLen") + 1;
		Logger log = Logger.getSingleton();
		// Allocate memory
		if (print) {
			log.debug("Model-3P d parameter: " + AlignmentModel
					.humanReadableDoubleCount(maxSrcSenLength * maxSrcSenLength
							* maxTrgtSenLength));
			log.debug("Model-3P o parameter: " + AlignmentModel
					.humanReadableDoubleCount(srcDictSize * maxWordLength));
		}
		d = new double[maxSrcSenLength]
						[maxSrcSenLength][maxTrgtSenLength];
		o = new double[srcDictSize][maxWordLength];
	}
	
	/**
	 * t, n, d are written to the file system.
	 * {@inheritDoc}
	 */
	@Override
	public void dumpToFilesystem(String prefix, String postfix) {
		Dictionary srcDict = Dictionary.getSingleton(Dictionary.SRC_DICT);
		Dictionary trgtDict = Dictionary.getSingleton(Dictionary.TRGT_DICT);
		this.dumpParameter(prefix + "n" + postfix, n);
		this.dumpParameter(prefix + "actual.n" + postfix, n, new Dictionary[]{
			srcDict, null});
		this.dumpParameter(prefix + "t" + postfix, t);
		this.dumpParameter(prefix + "actual.t" + postfix, t, new Dictionary[]{
			srcDict, null, trgtDict});
		this.dumpParameter(prefix + "d" + postfix, d);
		this.dumpParameter(prefix + "o" + postfix, o);
		this.dumpParameter(prefix + "actual.o" + postfix, o, new Dictionary[]{
			srcDict, null, null});
	}
}
